from flask import Flask, request, jsonify
from gmssl import sm2, func
from datetime import datetime, timedelta
import random

app = Flask(__name__)

class SM2CertificateAuthority:
    def __init__(self):
        # 生成国密SM2密钥对
        self.ca_key = sm2.CryptSM2.generate_private_key()
        self.ca_public_key = sm2.CryptSM2().load_public_key(self.ca_key)
        
        # 创建自签名根证书
        self.root_cert = self._create_root_certificate()
    
    def _create_root_certificate(self):
        """创建国密根证书（自签名）"""
        subject = {
            'C': 'CN',
            'O': '国密CA认证中心',
            'CN': '国密根证书颁发机构'
        }
        
        # 生成证书请求作为基础
        cert_req = func.new_certificate_request(
            subject=subject,
            public_key=self.ca_public_key
        )
        
        # 构建证书
        cert = func.new_x509_certificate(
            subject=subject,
            issuer=subject,
            public_key=self.ca_public_key,
            not_before=datetime.now(),
            not_after=datetime.now() + timedelta(days=3650),  # 10年有效期
            serial_number=random.randint(1, 2**64),
            signing_algorithm='sm2sign',
            sign_key=self.ca_key
        )
        
        # 添加扩展项
        cert = func.add_x509_extension(
            cert,
            'basicConstraints',
            'CA:TRUE',
            critical=True
        )
        cert = func.add_x509_extension(
            cert,
            'keyUsage',
            'keyCertSign, cRLSign',
            critical=True
        )
        
        return cert
    
    def issue_certificate(self, csr_pem, validity_days=365):
        """颁发国密证书"""
        # 加载CSR
        csr = func.load_certificate_request(csr_pem)
        
        # 验证CSR签名
        if not func.verify_certificate_request(csr):
            raise ValueError("无效的CSR签名")
        
        # 创建新证书
        cert = func.new_x509_certificate(
            subject=func.get_certificate_subject(csr),
            issuer=func.get_certificate_subject(self.root_cert),
            public_key=func.get_certificate_public_key(csr),
            not_before=datetime.now(),
            not_after=datetime.now() + timedelta(days=validity_days),
            serial_number=random.randint(1, 2**64),
            signing_algorithm='sm2sign',
            sign_key=self.ca_key
        )
        
        # 添加扩展项
        cert = func.add_x509_extension(
            cert,
            'keyUsage',
            'digitalSignature, keyEncipherment',
            critical=True
        )
        
        return func.dump_certificate(cert, encoding='PEM')

# 初始化国密CA
ca = SM2CertificateAuthority()

@app.route('/ca_root', methods=['GET'])
def get_ca_root():
    """获取CA根证书(PEM格式)"""
    return func.dump_certificate(ca.root_cert, encoding='PEM'), 200, {
        'Content-Type': 'application/x-pem-file'
    }

@app.route('/generate_csr', methods=['POST'])
def generate_csr():
    """生成国密证书签名请求(CSR)"""
    data = request.json
    subject = {
        'C': data.get('country', 'CN'),
        'ST': data.get('state', 'Beijing'),
        'L': data.get('locality', 'Beijing'),
        'O': data.get('organization', 'Test Organization'),
        'CN': data.get('common_name', 'example.com')
    }
    
    # 生成密钥对
    private_key = sm2.CryptSM2.generate_private_key()
    public_key = sm2.CryptSM2().load_public_key(private_key)
    
    # 创建CSR
    csr = func.new_certificate_request(
        subject=subject,
        public_key=public_key,
        signing_algorithm='sm2sign',
        sign_key=private_key
    )
    
    return jsonify({
        'private_key': private_key,
        'csr': func.dump_certificate_request(csr, encoding='PEM')
    }), 200

@app.route('/request_cert', methods=['POST'])
def request_certificate():
    """颁发国密证书"""
    data = request.json
    if not data or 'csr' not in data:
        return jsonify({'error': '请求中缺少CSR'}), 400
    
    try:
        certificate = ca.issue_certificate(
            data['csr'],
            validity_days=data.get('validity_days', 365)
        return jsonify({'certificate': certificate}), 200
    except Exception as e:
        return jsonify({'error': str(e)}), 400

@app.route('/verify_cert', methods=['POST'])
def verify_certificate():
    """验证国密证书"""
    data = request.json
    if not data or 'certificate' not in data:
        return jsonify({'error': '缺少证书参数'}), 400
    
    try:
        cert = func.load_certificate(data['certificate'])
        # 验证证书签名
        is_valid = func.verify_certificate(cert, ca.root_cert)
        return jsonify({
            'valid': is_valid,
            'subject': func.get_certificate_subject(cert),
            'issuer': func.get_certificate_issuer(cert),
            'serial_number': func.get_certificate_serial_number(cert),
            'not_before': func.get_certificate_not_before(cert),
            'not_after': func.get_certificate_not_after(cert)
        }), 200
    except Exception as e:
        return jsonify({'error': str(e)}), 400

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=28291, ssl_context='adhoc')